#!/usr/bin/python3
# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt


def ComputeTrajectory(waypoints,
                      vlim=np.r_[-np.ones((1, 6)), np.ones((1, 6))],
                      alim=np.r_[-5.0 *
                                 np.ones((1, 6)), 5.0 * np.ones((1, 6))],
                      v0=0.0, v1=0.0, gridpoints_number=200):
    """Compute Toppra joint trajectory

    Arguments:
        waypoints {np.array} -- trajectory waypoints, dimension: waypoints_number x dof
        vlim {np.arry} -- joint velocity limits, dimension: 2 (min, max) x dof
        alim {np.array} -- joint acceleration limits, dimension: 2 (min, max) x dof

    Keyword Arguments:
        v0 {float} -- start speed (default: {0.0})
        v1 {float} -- final speed (default: {0.0})
        gridpoints_number {int} -- compution stages, the larger, the better the trajectory is, but large number cost more time to calculate (default: {200})

    Returns:
        toppra.interpolator.SplineInterpolator -- joint trajectory, which has methods of eval(t)/evald(t)/evaldd(t) to get position/velocity/acceleration at t
    """
    import toppra
    import toppra.constraint as constraint
    import toppra.algorithm as algo
    # Random waypoints used to obtain a random geometric path. Here,
    # we use spline interpolation.
    waypoints = np.asarray(waypoints)
    waypoints_number = len(waypoints)
    path = toppra.SplineInterpolator(
        np.linspace(0, 1, waypoints_number), waypoints)
    # Create acceleration bounds, then acceleration constraint object
    pc_vel = constraint.JointVelocityConstraint(vlim.T)
    pc_acc = constraint.JointAccelerationConstraint(
        alim.T, discretization_scheme=constraint.DiscretizationType.Interpolation)
    # Setup a parametrization instance, then retime
    gridpoints = np.linspace(0, path.duration, gridpoints_number)
    instance = algo.TOPPRA([pc_vel, pc_acc], path,
                           gridpoints=gridpoints, solver_wrapper='seidel')
    traj, _ = instance.compute_trajectory(v0, v1, return_data=False)
    return traj


def PlotTraj(traj, smaple_number=100):
    """Plot a toppra trajectory

    Arguments:
        traj {toppra.interpolator.SplineInterpolator} -- joint trajectory
    """
    ts_sample = np.linspace(0, traj.get_duration(), smaple_number)
    qdds_sample = traj.evaldd(ts_sample)
    qds_sample = traj.evald(ts_sample)
    qs_sample = traj.eval(ts_sample)
    fig, axs = plt.subplots(1, 3, sharex=True, figsize=[12, 4])
    for i in range(6):
        axs[0].plot(ts_sample, qdds_sample[:, i], label="J{:d}".format(i + 1))
        axs[1].plot(ts_sample, qds_sample[:, i], label="J{:d}".format(i + 1))
        axs[2].plot(ts_sample, qs_sample[:, i], label="J{:d}".format(i + 1))
    axs[0].legend()
    axs[0].set_xlabel("Time (s)")
    axs[0].set_ylabel("Joint acceleration (rad/s^2)")
    axs[1].legend()
    axs[1].set_xlabel("Time (s)")
    axs[1].set_ylabel("Joint velocity (rad/s)")
    axs[2].legend()
    axs[2].set_xlabel("Time (s)")
    axs[2].set_ylabel("Joint position (rad)")
    plt.show()


if __name__ == "__main__":
    dof = 6
    np.random.seed(0)
    way_pts = np.random.randn(5, dof)
    # Create velocity bounds
    vlim = np.array([
        [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0],
        [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
    ])
    # Create acceleration bounds
    alim = np.array([
        [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0],
        [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
    ]) * 5.0
    traj = ComputeTrajectory(way_pts, vlim, alim)
    PlotTraj(traj)
